import os
import cv2
import math
import json
import socket
import threading
import numpy as np


from datetime import datetime
from collections import deque, Counter
from ais_bench.infer.interface import InferSession
from tf_transformations import euler_from_quaternion


class Decider:
    def __init__(self, window_size=5, threshold=0.8, ):

        self.windows = [deque(maxlen=window_size), deque(maxlen=window_size)] # 图像分类窗口

        self.degree_queue = deque([0], maxlen=window_size)  # degree 缓存
        self.degree_range = [-30, -2, 2, 30] # 斜坡角度范围 [stop <- x0 -> slope <- x1 -> level <- x2 -> slop <- x3 -> stop]
        self.degree_info = 'None'

        self.threshold = threshold  # 置信度阈值

        self.stable_num = 4 # 窗口中的稳定值

        self.decision_map = {
            '00': 'level', # 两个都是平路
            '01': 'overleap',
            '02': 'overleap',
            '10': 'stairs',
            '11': 'stairs',
            '12': 'stairs',
            '20': 'stairs',
            '21': 'stairs',
            '22': 'stairs',
        }

        self.key = 'None'
        self.decision = 'None'
        self.status = 'None'
    
    def add_frame(self, pred_class, confidence):
        for i, pred in enumerate(pred_class):
            confd = confidence[i]
            # 筛选高置信度帧
            if confd >= self.threshold:
                self.windows[i].append(pred)

        self.update_status()

    def set_degree_info(self):
        mean_degree = sum(self.degree_queue) / len(self.degree_queue)
        
        if mean_degree <= self.degree_range[0]:
            return 'is_stop'
        elif self.degree_range[0] < mean_degree <= self.degree_range[1]:
            return 'is_slope'
        elif self.degree_range[1] < mean_degree <= self.degree_range[2]:
            return 'is_level'
        elif self.degree_range[2] < mean_degree <= self.degree_range[3]:
            return 'is_slope'
        elif self.degree_range[3] < mean_degree:
            return 'is_stop'
        else:
            return 'None'
        
    def update_status(self):
        key = ['N', 'N']
        for i, window in enumerate(self.windows):
            if len(window) == 0:
                continue
            counts = Counter(window)
            stable_class, count = counts.most_common(1)[0]
            if count >= self.stable_num:
                key[i] = str(stable_class)
        key = ''.join(key)
        
        self.key = key if 'N' not in key else f'{key}:None'
        self.decision = self.decision_map[key] if 'N' not in key else self.decision
        self.degree_info = self.set_degree_info()

        if self.status == 'stairs':
            if self.degree_info == 'is_level': 
                self.status = self.decision
            else:
                pass
        else:
            if self.degree_info == 'is_level': 
                self.status = self.decision
            elif self.degree_info == 'is_slope':
                self.status = 'slope'
            elif self.degree_info == 'is_stop':
                self.status = 'stop'


class RoadClass:
    def __init__(
        self, 
        om1_file='/home/Code/RoboticDog/resource/model/terrains1_bs1_512_aipp.om',
        om2_file='/home/Code/RoboticDog/resource/model/terrains2_bs1_512_aipp.om',
        save_img_path='/home/nzd/save_imgs'
    ):
        
        # args
        device_id = 0
     
        self.box_x = 420
        self.box_y = 0
        self.box_widht = 1080
        self.box_height = 1080
        self.resize = (512, 512)
                
        # init
        self.om_sess1 = InferSession(device_id, om1_file)
        self.om_sess2 = InferSession(device_id, om2_file)

        self.decider = Decider()
        
        # 创建套接字服务器线程
        self.socket_thread = threading.Thread(target=self.start_socket_server)
        self.socket_thread.daemon = True
        self.socket_thread.start()
        
        print('欧拉角发布器已启动')

        # debug
        self.debug = True
        self.save_img_path = save_img_path
        if not os.path.exists(self.save_img_path):
            os.makedirs(self.save_img_path)

    def preprocess(self, img):
        # 1. 提取框内图像
        cropped_img = img[self.box_y : self.box_y+self.box_height, self.box_x : self.box_x+self.box_widht]

        # 2. resize
        resized_img = cv2.resize(cropped_img, self.resize)

        return resized_img

    def start_socket_server(self):
        # 创建TCP套接字
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        
        # 服务器地址和端口
        server_address = ('192.168.1.100', 8844)
        
        # 绑定地址和端口
        server_socket.bind(server_address)
        
        # 开始监听，最大连接数为1
        server_socket.listen(100)
        print(f"服务器正在监听 {server_address}...")
        
        while True:
            # 等待客户端连接
            print("等待客户端连接...")
            client_socket, client_address = server_socket.accept()
            try:
                print(f"客户端 {client_address} 已连接")
                
                while True:
                    # 接收数据
                    data = client_socket.recv(1024)
                    if not data:
                        print("客户端已断开连接")
                        break
                    
                    # 解析JSON数据
                    try:
                        # 解析顶层JSON
                        data_dict = json.loads(data.decode('utf-8'))
                        
                        # 检查是否包含odometry字段
                        if 'odometry' not in data_dict:
                            print("JSON数据缺少odometry字段")
                            continue
                            
                        # 从odometry中获取四元数
                        odom_data = data_dict['odometry']
                        qx = odom_data.get('qx')
                        qy = odom_data.get('qy')
                        qz = odom_data.get('qz')
                        qw = odom_data.get('qw')
                        
                        # 检查四元数分量是否存在
                        if None in [qx, qy, qz, qw]:
                            print.error("四元数分量不完整")
                            continue
                        
                        # 转换四元数为欧拉角
                        roll, pitch, yaw = euler_from_quaternion([qx, qy, qz, qw])
                        
                        # 发布roll和pitch
                        # angles.append(pitch)
                        degree = pitch * (180 / math.pi)
                        # print(degree)
                        self.decider.degree_queue.append(degree)

                    except json.JSONDecodeError as e:
                        print(f"JSON解析错误: {e}, 原始数据: {data.decode('utf-8')}")

                    except Exception as e:
                        print(f"数据处理异常: {e}")

                    
            except Exception as e:
                print(f"通信异常: {e}")
            finally:
                # 关闭客户端连接
                client_socket.close()

    def postpreocess(self, logits, axis=1):
        # 1. 数值稳定性处理：减去最大值防止指数溢出
        max_logits = np.max(logits, axis=axis, keepdims=True)
        exp_logits = np.exp(logits - max_logits)
        
        # 2. 计算 softmax 分母（沿指定轴求和）
        sum_exp = np.sum(exp_logits, axis=axis, keepdims=True)
        
        # 3. 计算 softmax 概率
        probabilities = exp_logits / sum_exp

        # 4. 获取分类indx，和概率
        label = np.argmax(probabilities, axis=1)[0]  # 每行最大值的索引
        prob = round(np.max(probabilities, axis=1)[0], 4)  # 每行的最大值
        
        return label, prob

    def save_info(self, image, label1, prob1, label2, prob2, windows, key, decision, degree_queue, degree_info, status):
        mean_degree = sum(degree_queue) / len(degree_queue)
        print(f"pred: {label1}:{prob1}_{label2}:{prob2}\nwindows: {windows}\nkey: {key}\ndecision: {decision}\nmean_degree: {mean_degree}\ndegree_info: {degree_info}\nstatus: {status}")
        
        now = datetime.now()
        formatted_time = now.strftime("%Y%m%d%H%M%S%f")[:-3]
        cv2.imwrite(f"{self.save_img_path}/{formatted_time}_{label1}:{prob1}_{label2}:{prob2}.jpg", image)


    def infer(self, image):
        img = self.preprocess(image)
        
        pred1 = self.om_sess1.infer([img], custom_sizes=1000)[0]
        pred2 = self.om_sess2.infer([img], custom_sizes=1000)[0]

        label1, prob1 = self.postpreocess(pred1)
        label2, prob2 = self.postpreocess(pred2)
        
        self.decider.add_frame([label1, label2], [prob1, prob2])
                        
        if self.debug:
           self.save_info(image, label1, prob1, label2, prob2, self.decider.windows, self.decider.key, self.decider.decision, self.decider.degree_queue, self.decider.degree_info, self.decider.status)

        return self.decider.status



# 以下是测试脚本

def test_angle():
    rc = RoadClass()
    
    while True:
        print("==============================")
        degree_queue = rc.decider.degree_queue
        mean_degree = sum(degree_queue) / len(degree_queue)

        print(f"degree_queue={degree_queue}, mean_degree={mean_degree}; degree_info={rc.decider.set_degree_info()}")
        time.sleep(0.03)

def test_all():
    rc = RoadClass()

    img_files = glob("./images_0619/*.jpg")
    img_files.sort()
    img_num = len(img_files)

    t1 = time.time()
    for img_file in img_files:
        print(f"[main]{img_file}: ")

        img = cv2.imread(img_file)
        if img is None:
            img_num -= 1
            continue
        
        rc.infer(img)

    use_time = time.time() - t1

    print(f"[main]img_num={img_num}, use_time={use_time}, fps={img_num/use_time}")


if __name__ == '__main__':
    import time
    from glob import glob

    # test_angle()
    test_all()
